package tree.avlTree;

import jdk.nashorn.internal.ir.IfNode;

/**
 * @author czt
 * @version 1.0
 * @since 2025/1/3
 * 二叉搜索树在插入和删除时，节点可能失衡
 * 如果在插入和删除时通过旋转，始终让二叉搜索树保持平衡，称为自平衡的二叉搜索树
 * AVL是自平衡二叉搜索树的实现之一
 */
public class AVLTree {

    static class AVLNode {
        int key;
        Object value;
        AVLNode left;
        AVLNode right;
        int height = 1; // 节点高度

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right, int height) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.height = height;
        }
    }

    // 求节点高度
    private int height(AVLNode node) {
        return node == null ? 0 : node.height;
    }

    // 更新节点高度(删除，添加，旋转)
    private void updateHeight(AVLNode node) {
        node.height = Math.max(height(node.left), height(node.right)) + 1;
    }

    /**
     * 平衡因子 = 左子树高度-右子树高度
     * @param node 当前节点
     * @return 平衡因子
     * 如果返回-1,0,1说明树是平衡的
     * 如果返回的值>1,<-1说明树不平衡
     * >1表示左边太高
     * <-1表示右边太高
     * LL 左边更高，失衡节点的左孩子的左边更高 旋转一次(右旋)
     * LR 左边更高，失衡节点的左孩子的右边更高 旋转两次(子树先左旋整体后右旋)
     * RL 右边更高，失衡节点的右孩子的左边更高 旋转两次(子树先右旋整体后左旋)
     * RR 右边更高，失衡节点的右孩子的右边更高 旋转一次(左旋)
     */
    private int bf(AVLNode node) {
        return height(node.left) - height(node.right);
    }

    /**
     * 左旋-解决RR
     * @param node 要旋转的节点(失衡的节点)
     * @return 新的根节点
     */
    private AVLNode leftRotate(AVLNode node) {
        AVLNode right = node.right;
        AVLNode rLeft = right.left;
        if (rLeft != null) {
            node.right = rLeft;
        }
        right.left = node;
        updateHeight(node);
        updateHeight(right);
        return right;
    }

    /**
     * 右旋-解决LL
     * @param node 要旋转的节点(失衡的节点)
     * @return 新的根节点
     */
    private AVLNode rightRotate(AVLNode node) {
        // 要旋转节点的左孩子
        AVLNode left = node.left;
        // 要旋转节点左孩子的右孩子
        AVLNode lRight = left.right;

        if (lRight != null) { // 可以不加这个非空判断
            // 如果不为空，则让它成为待旋转节点的左孩子(换爹)
            node.left = lRight;
        }
        // 上位
        left.right = node;
        // 更新高度顺序不能变
        updateHeight(node);
        updateHeight(left);
        return left;
    }

    // 先左旋后右旋，解决LR
    private AVLNode leftRightRotate(AVLNode node) {
        node.left = leftRotate(node.left);
        return rightRotate(node);
    }

    // 先右旋后左旋，解决RL
    private AVLNode rightLeftRotate(AVLNode node) {
        node.right = rightRotate(node.right);
        return leftRotate(node);
    }

    // 检查节点是否平衡，重新平衡代码
    private AVLNode balance(AVLNode node) {
        if (node == null) {
            return null;
        }
        int bf = bf(node);
        if (bf > 1 && bf(node.left) >= 0) {
            // LL情况
            return rightRotate(node);
        } else if (bf > 1 && bf(node.left) < 0) {
            // LR
            return leftRightRotate(node);
        } else if (bf < -1 && bf(node.right) > 0) {
            // RL
            return rightLeftRotate(node);
        } else if (bf < -1 && bf(node.right) <= 0) {
            // LL
            return rightRotate(node);
        }
        return node;
    }

    AVLNode root;

    // 添加
    public void put(int key, Object value) {
        root = doPut(root, key, value);
    }

    private AVLNode doPut(AVLNode node, int key, Object value) {
        // 1.找到空位，创建新节点
        if (node == null) {
            return new AVLNode(key, value);
        }
        // 2.key已经存在，更新
        if (key == node.key) {
            node.value = value;
        }
        // 3.继续查找
        if (key < node.key) {
            node.left = doPut(node.left, key, value);
        } else {
            node.right = doPut(node.right, key, value);
        }
        updateHeight(node);
        return balance(node);
    }

    // 删除
    public void remove(int key) {
        root = doRemove(root, key);
    }

    private AVLNode doRemove(AVLNode node, int key) {
        // 1.node == null
        if (node == null) {
            return null;
        }
        // 2.没找到 key
        if (key < node.key) {
            node.left = doRemove(node.left, key);
        } else if (key > node.key) {
            node.right = doRemove(node.right, key);
        } else {
            // 3.找到key  没有孩子  有一个孩子  有两个孩子
            if (node.left == null && node.right == null) {
                return null;
            } else if (node.left == null) {
                node = node.right;
            } else if (node.right == null) {
                node = node.left;
            } else {
                AVLNode s = node.right;
                while (s.left != null) {
                    s = s.left;
                }
                // s 后继节点
                s.right = doRemove(node.right, s.key);
                s.left = node.left;
                node = s;
            }
        }
        // 4.更新高度
        updateHeight(node);
        // 5.balance
        return balance(node);
    }

}
